import os

import matplotlib.image as mpig
import numpy as np

data_dir = r"C:\Users\Administrator\Downloads\captcha_images"
fileList = os.listdir(data_dir)
char_dict = {
    "0": 0, "1": 1, "2": 2, "3": 3, "4": 4,
    "5": 5, "6": 6, "7": 7, "8": 8, "9": 9,
    "A": 10, "B": 11, "C": 12, "D": 13, "E": 14,
    "F": 15, "G": 16, "H": 17, "I": 18, "J": 19,
    "K": 20, "L": 21, "M": 22, "N": 23, "O": 24,
    "P": 25, "Q": 26, "R": 27, "S": 28, "T": 29,
    "U": 30, "V": 31, "W": 32, "X": 33, "Y": 34,
    "Z": 35
}


def load_captcha(normalize=True, flatten=False, one_hot_label=True, transpose=False):
    """

    读入数据集
    Returns
    -------
    (训练图像, 训练标签), (测试图像, 测试标签)
    """
    imgs = []
    ts = []

    for file_name in fileList:
        img = mpig.imread(f"{data_dir}/{file_name}")
        t = np.zeros((4, 36))
        for i in range(4):
            idx = char_dict[file_name[i]]
            t[i][idx] = 1

        if flatten:
            ts.append(t.flatten())
        else:
            ts.append(t)

        imgs.append(img)

    imgs = np.array(imgs)
    if not normalize:
        imgs *= 255

    ts = np.array(ts)
    if not one_hot_label:
        ts = np.argmax(ts, -1)

    if transpose:
        imgs = imgs.transpose((0, 3, 1, 2))

    dataset = {}
    img_size = imgs.shape[0]
    test_size = 1000
    test_mask = np.random.choice(img_size, test_size)
    train_mask = np.delete(np.arange(0, img_size), test_mask)
    dataset['train_img'] = imgs[train_mask]
    dataset['train_label'] = ts[train_mask]

    dataset['test_img'] = imgs[test_mask]
    dataset['test_label'] = ts[test_mask]
    return (dataset['train_img'], dataset['train_label']), (dataset['test_img'], dataset['test_label'])
